use anyhow::Result;
use rusqlite::{params, OptionalExtension};
use serde::{Deserialize, Serialize};
use tauri::State;
use aes_gcm::{Aes256Gcm, Key, Nonce, KeyInit};
use aes_gcm::aead::Aead;
use argon2::Argon2;
use base64::{Engine as _, engine::general_purpose};
use rand::Rng;
use zeroize::Zeroize;

use super::storage::AppDb;
use super::password::MasterPasswordCache;

/// Generate encryption key from master password
fn derive_encryption_key(master_password: &str) -> Result<[u8; 32], String> {
    let salt = b"ai_model_manager_salt"; // Fixed salt for key derivation
    let argon2 = Argon2::default();
    
    let mut key = [0u8; 32];
    argon2
        .hash_password_into(master_password.as_bytes(), salt, &mut key)
        .map_err(|e| format!("密钥派生失败: {}", e))?;
    
    Ok(key)
}

/// Encrypt data with master password
fn encrypt_data(data: &str, master_password: &str) -> Result<String, String> {
    let key_bytes = derive_encryption_key(master_password)?;
    let key = Key::<Aes256Gcm>::from_slice(&key_bytes);
    let cipher = Aes256Gcm::new(key);
    
    // Generate random nonce
    let mut nonce_bytes = [0u8; 12];
    rand::thread_rng().fill(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);
    
    // Encrypt data
    let encrypted = cipher
        .encrypt(nonce, data.as_bytes())
        .map_err(|e| format!("加密失败: {}", e))?;
    
    // Combine nonce and encrypted data
    let mut encrypted_data = nonce_bytes.to_vec();
    encrypted_data.extend_from_slice(&encrypted);
    
    // Encode as base64
    let encrypted_b64 = general_purpose::STANDARD.encode(&encrypted_data);
    
    // Clear sensitive data
    let mut key_bytes = key_bytes;
    key_bytes.zeroize();
    
    Ok(encrypted_b64)
}

/// Decrypt data with master password
pub fn decrypt_data(encrypted_b64: &str, master_password: &str) -> Result<String, String> {
    let key_bytes = derive_encryption_key(master_password)?;
    let key = Key::<Aes256Gcm>::from_slice(&key_bytes);
    let cipher = Aes256Gcm::new(key);
    
    // Decode from base64
    let encrypted_data_bytes = general_purpose::STANDARD
        .decode(encrypted_b64)
        .map_err(|e| format!("Base64解码失败: {}", e))?;
    
    if encrypted_data_bytes.len() < 12 {
        return Err("加密数据格式错误".to_string());
    }
    
    // Extract nonce and encrypted data
    let nonce_bytes = &encrypted_data_bytes[..12];
    let encrypted = &encrypted_data_bytes[12..];
    
    let nonce = Nonce::from_slice(nonce_bytes);
    
    // Decrypt data
    let decrypted = cipher
        .decrypt(nonce, encrypted)
        .map_err(|e| format!("解密失败: {}", e))?;
    
    let data = String::from_utf8(decrypted)
        .map_err(|e| format!("数据解码失败: {}", e))?;
    
    // Clear sensitive data
    let mut key_bytes = key_bytes;
    key_bytes.zeroize();
    
    Ok(data)
}

/// AI Model representation
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct AiModel {
    pub id: Option<i64>,
    pub name: String,
    pub model_id: String,
    pub model_type: String,
    pub api_endpoint: String,
    pub api_key: String, // Decrypted API key
    pub is_default: bool,
    pub is_enabled: bool,
    pub max_tokens: Option<i64>,
    pub temperature: Option<f64>,
    pub description: Option<String>,
    pub created_at: Option<String>,
    pub updated_at: Option<String>,
}

/// Create a new AI model
#[tauri::command]
pub async fn create_ai_model(
    db: State<'_, AppDb>,
    password_cache: State<'_, MasterPasswordCache>,
    name: String,
    model_id: String,
    model_type: String,
    api_endpoint: String,
    api_key: String,
    is_default: bool,
    max_tokens: Option<i64>,
    temperature: Option<f64>,
    description: Option<String>,
) -> Result<i64, String> {
    // Get master password from cache
    let cache = password_cache.0.lock().map_err(|e| e.to_string())?;
    let master_password = cache.as_ref()
        .ok_or_else(|| "未找到主密码，请重新登录".to_string())?;
    
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    // Encrypt API key
    let encrypted_key = encrypt_data(&api_key, &master_password)
        .map_err(|e| format!("Failed to encrypt API key: {}", e))?;
    
    // If this model is set as default, unset other defaults
    if is_default {
        conn.execute(
            "UPDATE ai_models SET is_default = 0 WHERE is_default = 1",
            [],
        )
        .map_err(|e| format!("Failed to update default models: {}", e))?;
    }
    
    // Insert new model
    conn.execute(
        "INSERT INTO ai_models (name, model_id, model_type, api_endpoint, encrypted_api_key, is_default, max_tokens, temperature, description) 
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
        params![
            name,
            model_id,
            model_type,
            api_endpoint,
            encrypted_key,
            if is_default { 1 } else { 0 },
            max_tokens,
            temperature,
            description,
        ],
    )
    .map_err(|e| format!("Failed to insert model: {}", e))?;
    
    Ok(conn.last_insert_rowid())
}

/// Get all AI models
#[tauri::command]
pub async fn get_ai_models(
    db: State<'_, AppDb>,
    password_cache: State<'_, MasterPasswordCache>,
) -> Result<Vec<AiModel>, String> {
    // Get master password from cache
    let cache = password_cache.0.lock().map_err(|e| e.to_string())?;
    let master_password = cache.as_ref()
        .ok_or_else(|| "未找到主密码，请重新登录".to_string())?;
    
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    let mut stmt = conn
        .prepare(
            "SELECT id, name, model_id, model_type, api_endpoint, encrypted_api_key, is_default, max_tokens, temperature, description, created_at, updated_at 
             FROM ai_models 
             ORDER BY is_default DESC, created_at DESC"
        )
        .map_err(|e| e.to_string())?;
    
    let models = stmt
        .query_map([], |row| {
            let encrypted_key: String = row.get(5)?;
            let api_key = decrypt_data(&encrypted_key, &master_password)
                .unwrap_or_else(|_| String::from("[Decryption failed]"));
            
            Ok(AiModel {
                id: row.get(0)?,
                name: row.get(1)?,
                model_id: row.get(2)?,
                model_type: row.get(3)?,
                api_endpoint: row.get(4)?,
                api_key,
                is_default: row.get::<_, i32>(6)? != 0,
                is_enabled: row.get::<_, i32>(7).unwrap_or(1) != 0,
                max_tokens: row.get(8)?,
                temperature: row.get(9)?,
                description: row.get(10)?,
                created_at: row.get(11)?,
                updated_at: row.get(12)?,
            })
        })
        .map_err(|e| e.to_string())?
        .collect::<rusqlite::Result<Vec<_>>>()
        .map_err(|e| e.to_string())?;
    
    Ok(models)
}

/// Get a specific AI model by ID
#[tauri::command]
pub async fn get_ai_model(
    db: State<'_, AppDb>,
    password_cache: State<'_, MasterPasswordCache>,
    id: i64,
) -> Result<AiModel, String> {
    // Get master password from cache
    let cache = password_cache.0.lock().map_err(|e| e.to_string())?;
    let master_password = cache.as_ref()
        .ok_or_else(|| "未找到主密码，请重新登录".to_string())?;
    
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    let model = conn
        .query_row(
            "SELECT id, name, model_id, model_type, api_endpoint, encrypted_api_key, is_default, is_enabled, max_tokens, temperature, description, created_at, updated_at 
             FROM ai_models 
             WHERE id = ?1",
            params![id],
            |row| {
                let encrypted_key: String = row.get(5)?;
                let api_key = decrypt_data(&encrypted_key, &master_password)
                    .unwrap_or_else(|_| String::from("[Decryption failed]"));
                
                Ok(AiModel {
                    id: row.get(0)?,
                    name: row.get(1)?,
                    model_id: row.get(2)?,
                    model_type: row.get(3)?,
                    api_endpoint: row.get(4)?,
                    api_key,
                    is_default: row.get::<_, i32>(6)? != 0,
                    is_enabled: row.get::<_, i32>(7).unwrap_or(1) != 0,
                    max_tokens: row.get(8)?,
                    temperature: row.get(9)?,
                    description: row.get(10)?,
                    created_at: row.get(11)?,
                    updated_at: row.get(12)?,
                })
            },
        )
        .map_err(|e| format!("Model not found: {}", e))?;
    
    Ok(model)
}

/// Update an AI model
#[tauri::command]
pub async fn update_ai_model(
    db: State<'_, AppDb>,
    password_cache: State<'_, MasterPasswordCache>,
    id: i64,
    name: String,
    model_id: String,
    model_type: String,
    api_endpoint: String,
    api_key: String,
    is_default: bool,
    max_tokens: Option<i64>,
    temperature: Option<f64>,
    description: Option<String>,
) -> Result<(), String> {
    // Get master password from cache
    let cache = password_cache.0.lock().map_err(|e| e.to_string())?;
    let master_password = cache.as_ref()
        .ok_or_else(|| "未找到主密码，请重新登录".to_string())?;
    
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    // Encrypt API key
    let encrypted_key = encrypt_data(&api_key, &master_password)
        .map_err(|e| format!("Failed to encrypt API key: {}", e))?;
    
    // If this model is set as default, unset other defaults
    if is_default {
        conn.execute(
            "UPDATE ai_models SET is_default = 0 WHERE is_default = 1 AND id != ?1",
            params![id],
        )
        .map_err(|e| format!("Failed to update default models: {}", e))?;
    }
    
    // Update model
    conn.execute(
        "UPDATE ai_models 
         SET name = ?1, model_id = ?2, model_type = ?3, api_endpoint = ?4, encrypted_api_key = ?5, 
             is_default = ?6, max_tokens = ?7, temperature = ?8, description = ?9, updated_at = CURRENT_TIMESTAMP 
         WHERE id = ?10",
        params![
            name,
            model_id,
            model_type,
            api_endpoint,
            encrypted_key,
            if is_default { 1 } else { 0 },
            max_tokens,
            temperature,
            description,
            id,
        ],
    )
    .map_err(|e| format!("Failed to update model: {}", e))?;
    
    Ok(())
}

/// Delete an AI model
#[tauri::command]
pub async fn delete_ai_model(
    db: State<'_, AppDb>,
    id: i64,
) -> Result<(), String> {
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    conn.execute(
        "DELETE FROM ai_models WHERE id = ?1",
        params![id],
    )
    .map_err(|e| format!("Failed to delete model: {}", e))?;
    
    Ok(())
}

/// Set a model as default
#[tauri::command]
pub async fn set_default_ai_model(
    db: State<'_, AppDb>,
    id: i64,
) -> Result<(), String> {
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    // Unset all defaults
    conn.execute(
        "UPDATE ai_models SET is_default = 0 WHERE is_default = 1",
        [],
    )
    .map_err(|e| format!("Failed to unset defaults: {}", e))?;
    
    // Set new default
    conn.execute(
        "UPDATE ai_models SET is_default = 1, updated_at = CURRENT_TIMESTAMP WHERE id = ?1",
        params![id],
    )
    .map_err(|e| format!("Failed to set default: {}", e))?;
    
    Ok(())
}

/// Get the default AI model
#[tauri::command]
pub async fn get_default_ai_model(
    db: State<'_, AppDb>,
    password_cache: State<'_, MasterPasswordCache>,
) -> Result<Option<AiModel>, String> {
    // Get master password from cache
    let cache = password_cache.0.lock().map_err(|e| e.to_string())?;
    let master_password = cache.as_ref()
        .ok_or_else(|| "未找到主密码，请重新登录".to_string())?;
    
    let conn = db.0.lock().map_err(|e| e.to_string())?;
    
    let model = conn
        .query_row(
            "SELECT id, name, model_id, model_type, api_endpoint, encrypted_api_key, is_default, is_enabled, max_tokens, temperature, description, created_at, updated_at 
             FROM ai_models 
             WHERE is_default = 1 
             LIMIT 1",
            [],
            |row| {
                let encrypted_key: String = row.get(5)?;
                let api_key = decrypt_data(&encrypted_key, &master_password)
                    .unwrap_or_else(|_| String::from("[Decryption failed]"));
                
                Ok(AiModel {
                    id: row.get(0)?,
                    name: row.get(1)?,
                    model_id: row.get(2)?,
                    model_type: row.get(3)?,
                    api_endpoint: row.get(4)?,
                    api_key,
                    is_default: row.get::<_, i32>(6)? != 0,
                    is_enabled: row.get::<_, i32>(7).unwrap_or(1) != 0,
                    max_tokens: row.get(8)?,
                    temperature: row.get(9)?,
                    description: row.get(10)?,
                    created_at: row.get(11)?,
                    updated_at: row.get(12)?,
                })
            },
        )
        .optional()
        .map_err(|e| format!("Failed to get default model: {}", e))?;
    
    Ok(model)
}
